# 导入必要的库和自定义函数
import pandas as pd
from _data_process import cumulative_to_monthly, monthly_to_quarterly, calculate_yoy
from _data_plot import plot_double_axis_with_background
from 2关联分析数据 import merge_series_with_common_dates, add_same_trend_column

# 读取原始数据文件
elec_data_cum = pd.read_excel(f"data/全社会用电量当月累计值.xlsx", index_col=0)  # 读取全社会用电量累计值数据
gdp_data = pd.read_excel(f"data/GDP_不变价_当季同比.xlsx", index_col=0)  # 读取GDP当季同比数据

# 计算用电量的季度同比
# 第一步：将累计值转换为月度当月值
monthly_data = cumulative_to_monthly(elec_data_cum)  # 在原始数据是累计值的时候用
# 第二步：将月度数据转换为季度数据  
quarterly_data = monthly_to_quarterly(monthly_data)
# 第三步：计算季度同比增速
elec_data_quarterly_yoy = calculate_yoy(quarterly_data)

# 合并GDP和用电量数据
# 将两个时间序列合并为DataFrame，保留共同日期及最新两期非共同日期的数据
gdp_elec_data = merge_series_with_common_dates(
    gdp_data.iloc[:,0],  # GDP数据的第一列
    elec_data_quarterly_yoy.iloc[:,0],  # 用电量同比数据的第一列
    "GDP当季同比",  # GDP列的命名
    "全社会用电量当季同比"  # 用电量列的命名
).iloc[5:,:]  # 去掉前5行数据（通常用于去除初期不稳定的数据）

# 添加趋势分析列
# 为DataFrame添加四列：第一列涨跌、第二列涨跌、是否同趋势以及index的月份
gdp_elec_data = add_same_trend_column(gdp_elec_data)

# 绘制双轴对比图
# 使用自定义函数绘制GDP与用电量的双轴趋势对比图，背景色表示同趋势情况
fig_gdp_elec = plot_double_axis_with_background(
    gdp_elec_data.iloc[:,:],  # 使用所有列的数据
    gdp_elec_data.columns[0],  # 第一个列名（GDP当季同比）
    gdp_elec_data.columns[1],  # 第二个列名（全社会用电量当季同比）
    title="GDP与用电量关联分析",  # 图表标题
)